#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/4/9
# @Author  : geekhch
# @Email   : geekhch@qq.com
# @Desc    : None

class Evaluation():
    "NER 模型评价"
    def __init__(self):
        self.num_pred_entities = 0
        self.num_gold_entities = 0
        self.num_accu_entities = 0

        self.num_total_index = 0
        self.num_accu_index = 0

    def add_sample(self, gold_tag_ids, pred_tag_ids, seq_len, dataClass, text=None):
        '''
        :param gold_tags: list
        :param pred_tags: list, its length equal with gold_tags's
        '''
        if not isinstance(gold_tag_ids, list):
            gold_tag_ids = gold_tag_ids.tolist()
        if not isinstance(pred_tag_ids, list):
            pred_tag_ids = pred_tag_ids.tolist()

        i = 0
        pred_entity_from = -1  # 当前位置是否在实体内

        self.num_total_index += seq_len.item()

        while i < seq_len:
            pred_tag = dataClass.tags[pred_tag_ids[i]]
            gold_tag = dataClass.tags[gold_tag_ids[i]]
            if(gold_tag_ids[i] == pred_tag_ids[i]):
                self.num_accu_index += 1
            if pred_tag.startswith('B-'):
                self.num_pred_entities += 1
                pred_entity_from = i
            elif pred_entity_from != -1:
                # 是否与前面的共同组成一个entity, 如果不是则重置起始标志
                cur_tag_type = dataClass.tags[pred_tag_ids[pred_entity_from]]
                if not (pred_tag.startswith('I-') and pred_tag[2:] == cur_tag_type[2:]):
                    # 实体结束， 判断是否与真实值相同
                    if gold_tag_ids[pred_entity_from:i] == pred_tag_ids[pred_entity_from:i]:
                        self.num_accu_entities += 1

                    pred_entity_from = -1

            if gold_tag.startswith('B-'):
                self.num_gold_entities += 1
                if pred_tag != gold_tag and text is not None:
                    print(f'recall failed: {i} {text[i:i+5]} {text}')
            i += 1

    def eval_scores(self):
        # precision、recall, accuracy, f1-score
        precision = self.num_accu_entities / self.num_pred_entities
        recall = self.num_accu_entities/self.num_gold_entities
        accuracy = self.num_accu_index/self.num_total_index
        f1 = 2*(precision*recall)/(precision+recall+1e-6)
        return {'precision':precision,'recall':recall, 'accuracy':accuracy, 'f1':f1}


